from flask import Flask, request, jsonify, send_file
import xlwt
import sqlite3
from datetime import datetime
import io
import os

app = Flask(__name__)

def query_data_by_date_range(start_date, end_date, table_name, db_path):
    """
    根据日期范围从数据库查询数据
    
    Args:
        start_date: 开始日期 (YYYY-MM-DD)
        end_date: 结束日期 (YYYY-MM-DD)
        table_name: 数据库表名
        db_path: 数据库路径
    
    Returns:
        查询结果列表
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # 假设表中有一个date字段用于日期筛选
    query = f"""
    SELECT * FROM {table_name} 
    WHERE date_column BETWEEN ? AND ?
    ORDER BY date_column
    """
    
    cursor.execute(query, (start_date, end_date))
    results = cursor.fetchall()
    
    # 获取列名
    column_names = [description[0] for description in cursor.description]
    
    conn.close()
    
    return results, column_names

def create_excel_from_data(data, headers, sheet_name="Sheet1"):
    """
    将数据写入Excel文件
    
    Args:
        data: 查询结果数据列表
        headers: 表头列表
        sheet_name: 工作表名称
    
    Returns:
        Excel文件的字节流
    """
    # 创建工作簿
    workbook = xlwt.Workbook(encoding='utf-8')
    worksheet = workbook.add_sheet(sheet_name)
    
    # 设置样式
    header_style = xlwt.XFStyle()
    font = xlwt.Font()
    font.bold = True
    header_style.font = font
    
    # 写入表头
    for col, header in enumerate(headers):
        worksheet.write(0, col, header, header_style)
    
    # 写入数据
    for row, record in enumerate(data, start=1):
        for col, value in enumerate(record):
            worksheet.write(row, col, str(value) if value is not None else "")
    
    # 自动调整列宽
    for col in range(len(headers)):
        max_width = max(
            [len(str(headers[col]))] + 
            [len(str(row[col])) if row[col] is not None else 0 for row in data]
        )
        worksheet.col(col).width = min(max_width * 256, 256 * 50)  # 最大宽度限制
    
    # 将工作簿保存到字节流
    output = io.BytesIO()
    workbook.save(output)
    output.seek(0)
    
    return output

@app.route('/export_excel', methods=['POST'])
def export_excel():
    """
    Excel导出接口
    接收开始日期、结束日期、表名等参数，返回Excel文件
    """
    try:
        # 获取请求参数
        data = request.get_json()
        
        start_date = data.get('start_date')
        end_date = data.get('end_date')
        table_name = data.get('table_name', 'default_table')
        db_path = data.get('db_path', 'database.db')
        
        # 验证日期格式
        try:
            datetime.strptime(start_date, '%Y-%m-%d')
            datetime.strptime(end_date, '%Y-%m-%d')
        except ValueError:
            return jsonify({'error': '日期格式错误，请使用YYYY-MM-DD格式'}), 400
        
        # 查询数据
        records, headers = query_data_by_date_range(start_date, end_date, table_name, db_path)
        
        if not records:
            return jsonify({'message': '在指定日期范围内未找到数据'}), 404
        
        # 生成Excel
        excel_stream = create_excel_from_data(records, headers, table_name)
        
        # 返回Excel文件
        filename = f"{table_name}_{start_date}_to_{end_date}.xls"
        return send_file(
            excel_stream,
            mimetype='application/vnd.ms-excel',
            as_attachment=True,
            download_name=filename
        )
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

# 创建示例数据库和表用于测试
def init_test_database():
    """初始化测试数据库"""
    conn = sqlite3.connect('test_database.db')
    cursor = conn.cursor()
    
    # 创建测试表
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS sales_data (
            id INTEGER PRIMARY KEY,
            date_column TEXT,
            product_name TEXT,
            quantity INTEGER,
            price REAL,
            total REAL
        )
    ''')
    
    # 插入测试数据
    test_data = [
        ('2023-01-15', '产品A', 10, 25.50, 255.00),
        ('2023-01-16', '产品B', 5, 40.00, 200.00),
        ('2023-01-17', '产品C', 8, 30.00, 240.00),
        ('2023-02-01', '产品A', 12, 25.50, 306.00),
        ('2023-02-02', '产品D', 3, 55.00, 165.00),
        ('2023-02-03', '产品B', 7, 40.00, 280.00),
    ]
    
    cursor.executemany(
        'INSERT OR REPLACE INTO sales_data (date_column, product_name, quantity, price, total) VALUES (?, ?, ?, ?, ?)',
        test_data
    )
    
    conn.commit()
    conn.close()

@app.route('/test_export', methods=['GET'])
def test_export():
    """测试导出功能的示例页面"""
    html_content = '''
    <!DOCTYPE html>
    <html>
    <head>
        <title>Excel导出测试</title>
        <style>
            body { font-family: Arial, sans-serif; margin: 40px; }
            .form-group { margin: 15px 0; }
            label { display: inline-block; width: 120px; }
            input, select { padding: 8px; margin: 5px; }
            button { padding: 10px 20px; background: #007bff; color: white; border: none; cursor: pointer; }
            button:hover { background: #0056b3; }
        </style>
    </head>
    <body>
        <h2>Excel导出测试</h2>
        <form id="exportForm">
            <div class="form-group">
                <label>开始日期:</label>
                <input type="date" id="startDate" name="start_date" value="2023-01-01" required>
            </div>
            <div class="form-group">
                <label>结束日期:</label>
                <input type="date" id="endDate" name="end_date" value="2023-12-31" required>
            </div>
            <div class="form-group">
                <label>表名:</label>
                <select id="tableName" name="table_name">
                    <option value="sales_data">销售数据</option>
                </select>
            </div>
            <div class="form-group">
                <label>数据库:</label>
                <select id="dbPath" name="db_path">
                    <option value="test_database.db">测试数据库</option>
                </select>
            </div>
            <button type="submit">导出Excel</button>
        </form>
        
        <script>
            document.getElementById('exportForm').addEventListener('submit', async function(e) {
                e.preventDefault();
                
                const formData = new FormData(this);
                const jsonData = {};
                for (let [key, value] of formData.entries()) {
                    jsonData[key] = value;
                }
                
                try {
                    const response = await fetch('/export_excel', {
                        method: 'POST',
                        headers: {
                            'Content-Type': 'application/json',
                        },
                        body: JSON.stringify(jsonData)
                    });
                    
                    if (response.ok) {
                        const blob = await response.blob();
                        const url = window.URL.createObjectURL(blob);
                        const a = document.createElement('a');
                        a.href = url;
                        a.download = response.headers.get('Content-Disposition')?.split('filename=')[1] || 'export.xls';
                        a.click();
                        window.URL.revokeObjectURL(url);
                    } else {
                        const error = await response.json();
                        alert('导出失败: ' + error.error);
                    }
                } catch (error) {
                    alert('请求失败: ' + error.message);
                }
            });
        </script>
    </body>
    </html>
    '''
    return html_content

if __name__ == '__main__':
    # 初始化测试数据库
    init_test_database()
    app.run(debug=True)